import datetime
import logging
import os
import base64
import json
from paddleocr import PaddleOCR
import re
import time
from werkzeug.utils import secure_filename

import torch
from torchvision import transforms
from PIL import Image

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 加密
aes_key = b'L03-Cumminschine'
iv_key = b'Cty-china-ER-app'

# 图片路径
image_file_path = "/tmp/"
model_path = r"models/model_draft_1129.pth"
loaded_model = None
ocr = PaddleOCR(
    use_angle_cls=True,
    cls_model_dir=r"models/cls",
    det_model_dir=r"models/det",
    rec_model_dir=r"models/rec"
)
# 角度不对 模糊 假的  未识别
class_labels = ['您拍摄照片的角度倾斜', '您拍摄的照片比较模糊', '您拍摄的照片可能存在翻拍', '']


# 上传文件
# type (1  ESN) (2 OEMName、ChassisNumber、VIN) (3 OEMName、VIN) (4、WarrantyStartDate、ChassisNumber) (5、WarrantyStartDate)
def upload(file_content, file_name, type):
    file_path = image_file_path + file_name
    logger.info(f"saving file into : {file_path}")
    __save_file(file_path, file_content)

    ocr_result = __run_model(file_path, type)
    msg = ''
    if ocr_result is None:
        msg = __single_cf_test(file_path)

    logger.info("识别内容" + str(ocr_result))

    __clear_file(file_path)

    return {"code": 1, "data": ocr_result, 'msg': msg}


def __run_model(img_path, type):
    result = ocr.ocr(img_path, cls=True)
    logger.info("ocr 结果解析前：" + str(result))
    if result and len(result) > 0:
        result_arr = [line[1][0] for line in result[0]]
        logger.info("图片文字：" + str(result_arr))
        return __resolve_ocr_result(result_arr, type)
    else:
        logger.warning("解析失败")
        return None


def __resolve_ocr_result(result_arr, type):
    ocr_results = {}
    # validate model outputs
    for i in range(len(result_arr)):
        v_vin = re.findall("[L]{1}[A-Z0-9]{16}", str(result_arr[i]))
        v_esn = re.findall("[7]{1}[0-9]{7}", str(result_arr[i]))
        v_chassis = re.findall("\d{4}-\d{2}-\d{2}", str(result_arr[i]))
        v_name = re.findall("\d{4}-\d{2}-\d{2}", str(result_arr[i]))
        v_date = re.findall("\d{4}-\d{2}-\d{2}", str(result_arr[i]))

        if v_vin and type in [2, 3]:
            ocr_results["VIN"] = v_vin[0]

        if v_name and type in [2, 3]:
            ocr_results["OEMName"] = v_name[0]

        if v_esn and type == 1:
            ocr_results["ESN"] = v_esn[0]

        if v_chassis and type in [2, 4]:
            ocr_results["ChassisNumber"] = v_chassis[0]

        if v_date and type in [4, 5]:
            ocr_results["WarrantyStartDate"] = v_date[0]

    return ocr_results


# 识别失败-报错处理
def __single_cf_test(image_path):
    global loaded_model
    image = Image.open(image_path).convert('RGB')
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    input_tensor = transform(image)
    input_batch = input_tensor.unsqueeze(0)
    if loaded_model is None:
        loaded_model = torch.load(model_path, map_location=torch.device('cpu'))
    loaded_model.eval()
    with torch.no_grad():
        output = loaded_model(input_batch)

    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    predicted_class = torch.argmax(probabilities).item()
    predicted_label = class_labels[predicted_class]

    return predicted_label


# save file to local path
def __save_file(file_path, file_content):
    try:
        with open(file_path, "wb") as lf:
            lf.write(file_content)
        return True
    except Exception as e:
        logger.error('成功保存失败发生错误：', str(e))
        return False


def __clear_file(file_path):
    try:
        os.remove(file_path)
    except Exception as e:
        logger.error('clear file error', str(e))


def lambda_handler(event, context):
    # logger.info(f"event: {event}")
    headers = event.get("headers")
    body = event.get('body')
    type = headers.get("type")
    file_name = headers.get("filename")

    if type is None:
        raise ValueError("type is invalid")
    if body is None:
        raise ValueError("body is invalid")
    if file_name is None:
        raise ValueError("filename is invalid")

    file_content = base64.b64decode(body.encode('utf-8'))

    result = upload(file_content, file_name, int(type))

    return {
        "statusCode": 200,
        "headers": {
            'Content-Type': 'application/json',
            'Access-Control-Allow-Origin': '*',
            'Access-Control-Allow-Headers': '*',
            'Access-Control-Allow-Methods': '*'
        },
        "body": json.dumps(result)
    }